import pickle
import numpy as np
import scipy.io.wavfile as wf
import python_speech_features as sf
import sklearn.svm as svm
import sklearn.metrics as sm
import sklearn.preprocessing as sp
import os
import librosa.display
import matplotlib.pyplot as plt
import seaborn as sns


# 整理样本
def search_files(directory):
    """
        检索目录下的所有wav文件 返回目录字典
        {“appple”:[url, url...],
        “kiwi”:[url, url...],....}
    """
    files_dict = {}
    for cur_dir, sub_dirs, files in os.walk(directory):  #cur_dir：当前目录，sub_dirs：下一级目录，file:音频文件
        for file in files:
            if file.endswith(".wav"):
                label = cur_dir.split(os.path.sep)[-1]
                if label not in files_dict:
                    files_dict[label] = []
                url = os.path.join(cur_dir, file)  #音频文件路径
                files_dict[label].append(url)   #files_dict：音频文件路径与label一一对应
    return files_dict


def files_mfc(file_urls):
    """
        读取语音文件 并且 MFCC化音频文件 梅尔频率倒谱系数
        返回整理好的 样本数据
    """
    x_data, y_data = [], []
    for label, urls in file_urls.items():  # 音频文件和label组成一个元组
        for file in urls:
            sample_rate, signs = wf.read(file)
            mfc = sf.mfcc(signs, sample_rate)
            # print(sample_rate)
            # ==>(1,13)
            sample = np.mean(mfc, axis=0)  # 求平均值
            x_data.append(sample)  # 音频特征数组
            y_data.append(label)  # 标签特征数组
    x_data = np.array(x_data)
    return x_data, y_data


def plot_confusion_matrix(y_true, y_pred, labels):
    cm = sm.confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
    plt.xlabel('Predicted labels')
    plt.ylabel('True labels')
    plt.title('Confusion Matrix')
    plt.show()


# 读取训练集 文件路径
train_urls = search_files("D:/Pycharm/trainvoice_depoyment/train")   #train_urls<——files_dict，音频文件路径与label一一对应
#print(train_urls)

# 整理训练集
train_x, train_y = files_mfc(train_urls)   #train_x, train_y为梅尔倒频谱处理后的训练数据，x为数据，y为标签
#print(train_x)

# 标签编码
encoder = sp.LabelEncoder()
train_y_label = encoder.fit_transform(train_y)  #处理标签
#print(train_x.shape, train_y_label.shape)

# 创建模型 SVM 并训练
model = svm.SVC(kernel="poly", degree=2, gamma="auto", probability=True)
model.fit(train_x, train_y_label)
file = open("D:/graduation design/SVM/model.pickle", "wb")    #使用了支持向量机（SVM）作为分类器，并将模型保存到了文件 "D:/Pycharm/trainvoice5/model.pickle" 中。
pickle.dump(model, file)
file.close()
# # 读取测试集 文件路径
# test_urls = search_files("D:/Pycharm/trainvoice5/test")
# #print(test_urls)
#
# # 整理测试集
# test_x, test_y = files_mfc(test_urls)
# test_y_label = encoder.fit_transform(test_y)
#
# # 预测测试集 输出分类报告
# prd_test_y = model.predict(test_x)
# print ('===============================================================')
# print(sm.classification_report(test_y_label, prd_test_y))
# print ('===============================================================')


def predict_audio_file(audio_file_path, model, encoder):
    # 读取音频文件并提取特征
    sample_rate, signs = wf.read(audio_file_path)
    mfc = sf.mfcc(signs, sample_rate)
    sample = np.mean(mfc, axis=0)  # 求平均值

    # 使用模型进行预测
    feature = np.array([sample])
    predicted_label = model.predict(feature)
    predicted_class = encoder.inverse_transform(predicted_label)  # 将编码的标签转换回原始标签

    return predicted_class[0]

# 调用 predict_audio_file 函数并输出结果
audio_file_path = "D:/Pycharm/trainvoice5/train/abnormal_treble/abnormal_treble+.wav"
predicted_result = predict_audio_file(audio_file_path, model, encoder)
print("预测结果:", predicted_result)